"""This module defines the Water Supply Dispatch model. The model is created using 
the pyoptinterface library.
"""

from optmodel.solver import get_solver  
from optmodel.solver import set_solver_parameters
from loguru import logger
from optmodel._model.hydro import AddHydroConstraints
from optmodel._model.objective import AddObjective
from optmodel.log_config import timer


def define_model(param : dict) -> object:
    """This function creates the model class depending on predefined solver.

    Parameters
    ----------
    params : dict
        parameters for the model

    Returns
    -------
    object
        A pyoptinterface Model object depending on the solver

    Raises
    ------
    ValueError
        Unsupported or undefined solver
    """
    solver = get_solver(param)
    model = solver.Model()
    model.param = param
    logger.info(f"成功创建模型，使用求解器 ")
    set_solver_parameters(model)
    logger.info(f"成功设置求解器参数")

    return model

def define_sets(model : object) -> None:
    """
    This function defines the sets for the model.

    Parameters
    ----------
    model : object
        A pyoptinterface Model object
    """
    param = model.param

    ########### 设置时间、水库、用户集合 ################
    basic_sets = ['time', 'reservoir', 'water_user']
    for set_name in basic_sets:
        setattr(model, set_name, param[set_name])

    ########### 设置外部水源集合（如桓仁）——这里只有供出水任务的水库，也包括集团内部的水库 ##############
    ext_sources = param.get('external_water_sources')
    if isinstance(ext_sources, dict):
        external_water_sources_set = list(ext_sources.keys())
    elif ext_sources is None:
        external_water_sources_set = []
    else:
        external_water_sources_set = list(ext_sources)
    model.external_water_sources = external_water_sources_set
    
    ########### 设置时间序列（时间点），包括最后一个时间点的下一个月 ################
    time_p = param['time']
    last_time = time_p[-1]
    year, month = map(int, last_time.split('-'))
    # 计算下一个月
    if month == 12:
        next_year = year + 1
        next_month = 1
    else:
        next_year = year
        next_month = month + 1
    next_time = f"{next_year:04d}-{next_month:02d}"
    model.time_p = time_p + [next_time]
    
    # 设置时间段序列（除第一个时间点外的所有时间点）
    model.time_nofirst = model.time_p[1:]

    ########### 设置使用弃水上限约束的水库集合（仅包含已在 param['spillflow_up'] 中出现的水库，保持现有顺序）###########
    spillflow_up_param = param.get('spillflow_up', {})  # 若用户没设置弃水上限约束，模型忽略
    spillflow_up_reservoirs = {r for (r, _) in spillflow_up_param.keys()}
    model.spillflow_up_cons_set = [r for r in model.reservoir if r in spillflow_up_reservoirs]

    ########### 设置使用弃水下限约束的水库集合（仅包含已在 param['spillflow_down'] 中出现的水库，保持现有顺序）###########
    spillflow_down_param = param.get('spillflow_down', {})  # 若用户没设置弃水下限约束，模型忽略
    spillflow_down_reservoirs = {r for (r, _) in spillflow_down_param.keys()}
    model.spillflow_down_cons_set = [r for r in model.reservoir if r in spillflow_down_reservoirs]

    ########### 设置使用 import_export 上限约束集合（仅包含在 param['import_export_up'] 中出现的三元组）###########
    # 该集合用于只对已配置的 (source_stnm, target_reservoir_stnm, time) 施加约束
    import_export_up_param = param.get('import_export_up', {})  # 若用户没设置引调水上限约束，模型忽略
    import_export_up_triples = list(import_export_up_param.keys())  # 每个键均为 (source, target, time)
    if import_export_up_triples:
        # 按现有集合顺序进行排序，保证约束生成顺序稳定
        def _triple_order(tri):
            s, r, t = tri
            try:
                s_idx = model.external_water_sources.index(s)
            except ValueError:
                s_idx = len(model.external_water_sources)
            try:
                r_idx = model.reservoir.index(r)
            except ValueError:
                r_idx = len(model.reservoir)
            try:
                t_idx = model.time.index(t)
            except ValueError:
                t_idx = len(model.time)
            return (s_idx, r_idx, t_idx)
        import_export_up_triples.sort(key=_triple_order)
    model.import_export_up_cons_set = import_export_up_triples

    ########### 设置使用 import_export 下限约束集合（仅包含在 param['import_export_down'] 中出现的三元组）###########
    # 该集合用于只对已配置的 (source_stnm, target_reservoir_stnm, time) 施加约束
    import_export_down_param = param.get('import_export_down', {})  # 若用户没设置引调水下限约束，模型忽略
    import_export_down_triples = list(import_export_down_param.keys())  # 每个键均为 (source, target, time)
    if import_export_down_triples:
        # 按现有集合顺序进行排序，保证约束生成顺序稳定
        import_export_down_triples.sort(key=_triple_order)
    model.import_export_down_cons_set = import_export_down_triples
    
    ########### 设置使用汛末库容约束的水库集合（仅包含已在 param['end_storage_flood_season_cons'] 中出现的水库）###########
    end_storage_flood_season_param = param.get('end_storage_flood_season_cons', {})  # 若用户没设置汛末库容约束，模型忽略
    end_storage_flood_season_reservoirs = list(end_storage_flood_season_param.keys())
    model.end_storage_flood_season_cons_set = [r for r in model.reservoir if r in end_storage_flood_season_reservoirs]
    
    ########### 设置使用调度期末库容约束的水库集合（仅包含已在 param['end_storage_dispatch_term_cons'] 中出现的水库）###########
    end_storage_dispatch_term_param = param.get('end_storage_dispatch_term_cons', {})  # 若用户没设置调度期末库容约束，模型忽略
    end_storage_dispatch_term_reservoirs = list(end_storage_dispatch_term_param.keys())
    model.end_storage_dispatch_term_cons_set = [r for r in model.reservoir if r in end_storage_dispatch_term_reservoirs]
    
def define_variables(model : object) -> None:
    """Define variables for the model.

    Parameters
    ----------
    model : object
        Model to be solved.
    """
    ########### 定义逐时入库量 ################
    model.inflow = model.add_variables(model.reservoir, model.time, lb=0)
    ########### 定义逐时段出库量 ################
    model.outflow = model.add_variables(model.reservoir, model.time, lb=0)
    ########### 定义逐时段供水流量 ################
    model.supplyflow = model.add_variables(model.reservoir, model.water_user, model.time, lb=0)
    ########### 定义水库库容 ################
    model.storage = model.add_variables(model.reservoir, model.time_p, lb=0)
    ########### 定义水库的逐时段弃水流量 ################
    model.spillflow = model.add_variables(model.reservoir, model.time, lb=0)
    ########### 定义引入水量 ################
    model.importflow = model.add_variables(model.external_water_sources, model.reservoir, model.time, lb=0)
    # model.importflow = model.add_variables(model.reservoir, model.time, lb=0)

def define_constraints(model : object) -> None:
    """Define constraints for the model.
    
    Parameters
    ----------
    model : object
        Model to be solved.
    """
    AddHydroConstraints(model)
    

def define_objective(model : object) -> None:
    """Add the  objective function to the model.

    Parameters
    ----------
    model : object
        Model to be solved.
    """
    AddObjective(model)

@timer
def creat_model(param : dict) -> object:
    model = define_model(param)
    define_sets(model)
    define_variables(model)
    define_constraints(model)
    define_objective(model)

    return model


if __name__ == '__main__':
    # 演示：从 set_up 加载 para，用 load_data 处理为 params，并创建模型
    import os
    import sys

    # 确保项目根目录在路径中，以便相对路径加载配置
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, project_root)

    from optmodel.set_up import load_water_supply_model
    from optmodel.load_data import process_para_data

    try:
        # 1) 加载原始数据 para（配置 + 径流 + 需水）
        config_path = os.path.join(project_root, "config.json")
        para = load_water_supply_model(config_path)

        # 2) 处理为模型参数 param
        param = process_para_data(para)
        logger.info(f"成功加载并处理配置文件 {config_path}")

        # 4) 创建模型
        model = creat_model(param)
        logger.info("bingo")

    except Exception as e:
        logger.error(f"示例运行失败: {e}")
